from typing import Optional, Callable

import pyorio.record.record as pyorio_record
from coin.support.data_replay.reader.interface import (
    RawData,
    IRawDataReader)


class SimpleReader(IRawDataReader):
  def __init__(self,
               start_timestamp: int,
               end_timestamp: int,
               file: str,
               timestamp_func: Callable[[bytes], None]):
    self._start_timestamp = start_timestamp
    self._end_timestamp = end_timestamp
    self._reader = pyorio_record.SimpleReader(file)
    self._timestamp_func = timestamp_func
    self._record_cache = None

  def peek(self) -> Optional[RawData]:
    ret = self._record_cache
    if ret is None:
      while True:
        record = self._reader.Read()
        if record is None:
          break
        timestamp = self._timestamp_func(record)
        if timestamp < self._start_timestamp:
          continue
        if timestamp > self._end_timestamp:
          break
        ret = RawData(timestamp=timestamp, value=record)
        self._record_cache = ret
        break
    return ret

  def read(self) -> Optional[RawData]:
    record = self._record_cache
    if record is None:
      record = self.peek()
    self._record_cache = None
    return record
